# -*- coding: utf-8 -*-
"""
Author: Qwen (基于 chen xi 脚本增强)
File: udp_stress_test_ramp_up.py
功能：支持平滑启动的 UDP 压力测试
      - 300秒内平滑启动200线程
      - 启动完成后持续运行（按时间或请求数）
      - 统计 QPS、延迟、成功率
"""

import struct
import socket
import time
import random
from datetime import datetime
import threading
from collections import deque
import statistics

# ==================== 配置参数 ====================
SERVER_IP = "120.26.136.5"
SERVER_PORT = 50001
TIMEOUT = 5  # UDP 超时（秒）

# 🚀 并发控制
MAX_THREADS = 100            # 最大线程数
RAMP_UP_TIME = 300           # 平滑启动时间（秒），300秒内启动200线程

# 📊 测试模式
DURATION_MODE = True         # True: 按时间运行；False: 按请求数运行
RUN_DURATION = 24 * 60 * 60  # 24 小时（仅 DURATION_MODE=True 时生效）
TOTAL_REQUESTS = 10_000_000  # 总请求数（仅 DURATION_MODE=False 时生效）

# 🧬 DID 范围
DID_PREFIX = "IOTGDD-"
DID_SUFFIX = "-JRIHMV"
MIN_INDEX = 0
MAX_INDEX = 60000  # 包含

# 📡 UDP 协议常量
START_CODE = 0xFFFFFFA1
MSG_TYPE = 3
MSG_HEADER_FORMAT = ">III"  # 大端序

# 🔢 统计锁
stats_lock = threading.Lock()

# 📈 全局统计
total_requests = 0
total_success = 0
total_timeout = 0
total_error = 0
latencies = []  # 存储延迟（秒）
start_time = None

# 🕒 QPS 滑动窗口（最近1秒）
qps_window = deque()
WINDOW_SIZE = 1  # 1秒


# ==================== 工具函数 ====================
def create_request(did):
    """构建 UDP 请求数据包"""
    payload = did.encode('utf-8')
    msg_len = len(payload)
    header = struct.pack(MSG_HEADER_FORMAT, START_CODE, MSG_TYPE, msg_len)
    return header + payload


def generate_random_did():
    """随机生成一个 DID"""
    index = random.randint(MIN_INDEX, MAX_INDEX)
    return f"{DID_PREFIX}{index:06d}{DID_SUFFIX}"


def worker():
    """工作线程函数：持续发送请求"""
    global total_requests, total_success, total_timeout, total_error, latencies

    while True:
        # 检查是否已达到总请求数限制（仅在按请求数模式下）
        if not DURATION_MODE:
            with stats_lock:
                if total_requests >= TOTAL_REQUESTS:
                    return

        # 检查是否已超时（按时间模式）
        if DURATION_MODE and (time.time() - start_time) >= RUN_DURATION:
            return

        # 生成 DID 并发送
        did = generate_random_did()
        request_data = create_request(did)
        sent_time = time.time()

        try:
            with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
                sock.settimeout(TIMEOUT)
                sock.sendto(request_data, (SERVER_IP, SERVER_PORT))
                response, _ = sock.recvfrom(1024)
            rtt = time.time() - sent_time

            with stats_lock:
                total_requests += 1
                total_success += 1
                latencies.append(rtt)
                qps_window.append(sent_time)

        except socket.timeout:
            rtt = time.time() - sent_time
            with stats_lock:
                total_requests += 1
                total_timeout += 1
                qps_window.append(sent_time)

        except Exception as e:
            rtt = time.time() - sent_time
            with stats_lock:
                total_requests += 1
                total_error += 1
                qps_window.append(sent_time)


def get_current_qps():
    """计算最近1秒的QPS"""
    now = time.time()
    while qps_window and qps_window[0] < now - WINDOW_SIZE:
        qps_window.popleft()
    return len(qps_window)


def print_stats_periodically():
    """每10秒打印一次实时统计"""
    global start_time
    while True:
        time.sleep(10)
        with stats_lock:
            elapsed = time.time() - start_time
            if total_requests == 0:
                continue
            qps = get_current_qps()
            avg_latency = statistics.mean(latencies[-1000:]) * 1000 if latencies else 0
            p95 = statistics.quantiles(latencies[-1000:], n=20)[-1] * 1000 if len(latencies) >= 1000 else 0
            success_rate = (total_success / total_requests) * 100

        print(f"[{datetime.now().strftime('%H:%M:%S')}] "
              f"QPS: {qps:6.1f} | "
              f"Success: {total_success:8,} | "
              f"Timeout: {total_timeout:6,} | "
              f"Error: {total_error:6,} | "
              f"Avg RTT: {avg_latency:6.2f}ms | "
              f"P95 RTT: {p95:6.2f}ms | "
              f"Duration: {int(elapsed//60)}m{int(elapsed%60)}s")


# ==================== 主逻辑 ====================
def main():
    global start_time

    print(f"🚀 UDP 压力测试启动（支持平滑启动）")
    print(f"目标：{RAMP_UP_TIME} 秒内平滑启动 {MAX_THREADS} 个线程")
    print(f"服务器: {SERVER_IP}:{SERVER_PORT}")
    print(f"最大线程数: {MAX_THREADS}")
    print(f"DID 范围: {MIN_INDEX} ~ {MAX_INDEX}")
    if DURATION_MODE:
        print(f"运行模式: 时间模式，持续 {RUN_DURATION//3600} 小时")
    else:
        print(f"运行模式: 请求数模式，共 {TOTAL_REQUESTS:,} 个请求")
    print("开始平滑启动线程...")

    # 启动统计线程
    stats_thread = threading.Thread(target=print_stats_periodically, daemon=True)
    stats_thread.start()

    start_time = time.time()
    threads = []

    # 🔁 平滑启动：在 RAMP_UP_TIME 秒内启动 MAX_THREADS 个线程
    for i in range(MAX_THREADS):
        t = threading.Thread(target=worker, daemon=True)
        threads.append(t)
        t.start()

        # 计算 sleep 时间，使线程均匀分布在 RAMP_UP_TIME 内
        if i < MAX_THREADS - 1:  # 最后一个线程不需要 sleep
            time.sleep(RAMP_UP_TIME / MAX_THREADS)

    print(f"✅ {MAX_THREADS} 个线程已全部启动，进入持续运行阶段...")

    # 等待测试结束
    if DURATION_MODE:
        time.sleep(max(0, RUN_DURATION - RAMP_UP_TIME))
    else:
        while total_requests < TOTAL_REQUESTS:
            time.sleep(1)

    # 测试结束，等待线程自然退出（worker 会自动判断结束条件）
    print("\n⏹️  测试结束，等待线程退出...")

    # 最终统计
    final_elapsed = time.time() - start_time
    overall_qps = total_requests / final_elapsed if final_elapsed > 0 else 0
    success_rate = (total_success / total_requests * 100) if total_requests > 0 else 0
    avg_latency_ms = (statistics.mean(latencies) * 1000) if latencies else 0
    p95_latency_ms = (statistics.quantiles(latencies, n=100)[89] * 1000) if len(latencies) >= 100 else 0
    p99_latency_ms = (statistics.quantiles(latencies, n=100)[98] * 1000) if len(latencies) >= 100 else 0

    print("\n" + "=" * 70)
    print("✅ 压力测试完成！")
    print(f"总运行时间     : {final_elapsed//3600:.0f}h {(final_elapsed%3600)//60:.0f}m {final_elapsed%60:.0f}s")
    print(f"总请求数       : {total_requests:,}")
    print(f"成功响应       : {total_success:,} ({success_rate:.2f}%)")
    print(f"超时           : {total_timeout:,}")
    print(f"其他错误       : {total_error:,}")
    print(f"平均 QPS       : {overall_qps:.2f}")
    print(f"平均延迟       : {avg_latency_ms:.2f} ms")
    print(f"P95 延迟       : {p95_latency_ms:.2f} ms")
    print(f"P99 延迟       : {p99_latency_ms:.2f} ms")
    print("=" * 70)


if __name__ == "__main__":
    main()